import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.statespace.sarimax import SARIMAX
import pmdarima as pm

# 设置字体以支持 Unicode 字符
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 读取数据
data = pd.read_excel("data/train.xlsx")
data = data.fillna(0)
data = data.dropna()

data = data['WIND'].astype(np.float32)

# 检查数据的基本统计信息
print(data.describe())

# 自动选择SARIMA模型参数
auto_model = pm.auto_arima(data, seasonal=True, m=12, trace=True, error_action='ignore', suppress_warnings=True)
print(f"自动选择的模型参数: {auto_model.order} 和季节性参数: {auto_model.seasonal_order}")

# 构建SARIMA模型
sarima_model = SARIMAX(data, order=auto_model.order, seasonal_order=auto_model.seasonal_order)
sarima_result = sarima_model.fit(disp=False)

# 预测未来100天的数据
n_forecast = 100
forecast = sarima_result.get_forecast(steps=n_forecast)
forecast_mean = forecast.predicted_mean
forecast_conf_int = forecast.conf_int()

# 绘制结果
plt.figure(figsize=(12, 6))
plt.plot(data, label='历史数据')
plt.plot(forecast_mean, label='预测值', color='red')
plt.fill_between(forecast_conf_int.index,
                 forecast_conf_int.iloc[:, 0],
                 forecast_conf_int.iloc[:, 1], color='pink', alpha=0.3, label='置信区间')
plt.legend()
plt.xlabel('时间')
plt.ylabel('销量(千克)')
plt.title('SARIMA模型预测未来100天的数据')
plt.show()

# 打印预测结果
print("未来100天的预测值:")
print(forecast_mean)

# 绘制残差图
residuals = sarima_result.resid
plt.figure(figsize=(12, 6))
plt.plot(residuals, label='残差')
plt.axhline(y=0, color='r', linestyle='--')
plt.legend()
plt.xlabel('时间')
plt.ylabel('残差')
plt.title('SARIMA模型残差图')
plt.show()